from data_visualization.data_visualization import DataVisualization


class ComboPltChart(DataVisualization):
    use_2nd_axis = True
    y_label_1 = ""
    y_values_1 = []
    y_label_2 = ""
    y_values_2 = []
    chart_type_1 = ""
    chart_type_2 = ""

    def __init__(self):
        DataVisualization.__init__(self)
        self.algorithm_name = "复合Matplotlib图表"
        self.ipynb_template_name = "combo_plt_chart-template.ipynb"
        
    def setYLabels(self, y_label_1, y_label_2):
        self.y_label_1 = y_label_1
        self.y_values_1 = self.data_source[self.y_label_1].values
        self.y_label_2 = y_label_2
        self.y_values_2 = self.data_source[self.y_label_2].values
        
    def draw(self, x, y, ax, chart_type, color):
        if chart_type == "bar":
            ax.bar(x, y, color=color)
        elif chart_type == "plot":
            ax.plot(x, y, color=color)
        
    def implent(self): 
        DataVisualization.implent(self)
        
        ax = self.figure.add_subplot(111)
        ax.set_title(self.chart_title)
        ax.set_xlabel(self.x_label)
        ax2 = ax
        if self.use_2nd_axis:
            ax2 = ax.twinx()
            ax.set_ylabel(self.y_label_1)
            ax2.set_ylabel(self.y_label_2)
            
        self.draw(self.x_values, self.y_values_1, ax, self.chart_type_1, color="b")
        self.draw(self.x_values, self.y_values_2, ax2, self.chart_type_2, color="r")
        
        self.figure.legend([self.y_label_1, self.y_label_2])
        self.saveToPng()
        
        if self.use_2nd_axis:
            self.code = self.code + """
fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_title("#chart_title#")
ax.#chart_type_1#(x_values, y_values_1)
ax.set_xlabel("#x_label#")
ax.set_ylabel("#y_label_1#")
ax2 = ax.twinx()
ax2.#chart_type_2#(x_values, y_values_2, color="r")
ax2.set_xlabel("#x_label#")
ax2.set_ylabel("#y_label_2#")
fig.legend(["#y_label_1#", "#y_label_2#"])
plt.show()
            """
        else:
            self.code = self.code + """
plt.title("#chart_title#")
plt.#chart_type_1#(x_values, y_values_1)
plt.#chart_type_2#(x_values, y_values_2)
plt.legend(["#y_label_1#", "#y_label_2#"])
plt.show()
            """

    def prepareIpynbItems(self):
        DataVisualization.prepareIpynbItems(self)
        self.ipynb_items["#y_label_1#"] = self.y_label_1
        self.ipynb_items["#y_values_1#"] = self.y_values_1
        self.ipynb_items["#y_label_2#"] = self.y_label_2
        self.ipynb_items["#y_values_2#"] = self.y_values_2
        self.ipynb_items["#chart_type_1#"] = self.chart_type_1
        self.ipynb_items["#chart_type_2#"] = self.chart_type_2
